package com.test.springai.service;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

/**
 * Created with IntelliJ IDEA.
 *
 * @description:
 * @author: liuziyang
 * @since: 9/3/25 11:06 AM
 * @modifiedBy:
 * @version: 1.0
 */
@Service
public class SimpleRagService {
  private final ChatClient chatClient;
  private final SimpleVectorStore vectorStore;

  private static final String SYSTEM_PROMPT =
      """
      你是公司内部相关制度流程的管理员，对于用户的使用需求作出解答。
      如果上下文中没有包含足够的信息来回答问题，请说明你无法回答，不要编造信息。
      上下文：{context}
    """;

  public SimpleRagService(ChatClient.Builder builder, SimpleVectorStore vectorStore) {
    this.chatClient = builder.build();
    this.vectorStore = vectorStore;
  }

  /**
   * 根据用户输入生成响应信息，该响应通过检索增强生成rag过程获得.
   *
   * @param userInput the user input
   * @return the string
   */
  public Flux<String> generateResponse(String userInput) {
    // 1.在SimpleVectorStore中进行相似度搜索，获取与用户输入相关的文档
    // vectorStore.similaritySearch方法可以直接接收字符串，它会在内部将用户输入转换成嵌入向量
    final List<Document> relevantDocuments = vectorStore.similaritySearch(userInput);
    // 2. 将检索到的文档内容提取并组合成一个连续的上下文字符串
    // 每个文档的内容之间用双换行符分隔，以提高模型的可读性
    final String context =
        relevantDocuments.stream().map(Document::getText).collect(Collectors.joining("\n\n"));
    // 3. 填充系统提示词模板，将检索到的上下文插入到预定义的系统提示中
    // 这将指导模型使用提供的上下文进行回答。
    var systemPromptTemplate = new SystemPromptTemplate(SYSTEM_PROMPT);
    // 使用Map.of将context变量映射到检索到的上下文内容中
    final SystemMessage systemMessage =
        new SystemMessage(systemPromptTemplate.createMessage(Map.of("context", context)).getText());
    // 4. 构建用户消息
    var userMessage = new UserMessage(userInput);
    // 5. 将包含上下文的系统消息和用户消息组合成一个完整的Prompt
    final Prompt prompt = new Prompt(List.of(systemMessage, userMessage));
    // 6. 调用ChatClient发送prompt给大模型，获取响应内容
    return chatClient.prompt(prompt).stream().content();
  }
}
